#include <Uefi.h>
#include <Library/UefiLib.h>
#include <Library/UefiBootServicesTableLib.h>

EFI_STATUS
TestMemoryAlloc()
{
  VOID *Buffer      = NULL;
  EFI_STATUS Status = EFI_SUCCESS;
  UINTN BufferSize  = 100;

  // Allocate memory
  Status = gBS->AllocatePool(EfiBootServicesData, BufferSize, (VOID **)&Buffer);
  if (EFI_ERROR(Status)) {
    Print(L"TestMemoryAlloc: AllocatePool failed: %llx\n", Status);
    return Status;
  }

  Print(L"TestMemoryAlloc: Memory allocated at address: %p\n", Buffer);

  // Use the allocated memory
  SetMem(Buffer, BufferSize, 0);

  // Free the allocated memory
  Status = gBS->FreePool(Buffer);
  if (EFI_ERROR(Status)) {
    Print(L"TestMemoryAlloc: FreePool failed: %llx\n", Status);
    return Status;
  }

  return EFI_SUCCESS;
}

// 在任意地址分配内存页
EFI_STATUS
TestAllocateAnyPages()
{
  EFI_STATUS Status           = EFI_SUCCESS;
  EFI_PHYSICAL_ADDRESS Memory = 0;
  UINTN Pages                 = 3; // Number of pages to allocate (1 page = 4KB)

  // Allocate pages
  Status = gBS->AllocatePages(AllocateAnyPages, EfiBootServicesData, Pages, &Memory);
  if (EFI_ERROR(Status)) {    
    Print(L"TestAllocateAnyPages: AllocatePages failed: %llx\n", Status);
    return Status;
  }

  Print(L"TestAllocateAnyPages: Pages allocated at address: %llx, Pages: %lu\n", Memory, Pages);

  // Use the allocated pages
  SetMem((VOID *)(UINTN)Memory, Pages * 4096, 0);

  // Free the allocated pages
  Status = gBS->FreePages(Memory, Pages);
  if (EFI_ERROR(Status)) {
    Print(L"TestAllocateAnyPages: FreePages failed: %llx\n", Status);
    return Status;
  }

  return EFI_SUCCESS;
}

// 在指定地址分配内存页
EFI_STATUS
TestAllocatePages()
{
  EFI_STATUS Status           = EFI_SUCCESS;
  EFI_PHYSICAL_ADDRESS Memory = 1024 * 1024 * 10; // Desired address (10MB)
  UINTN Pages                 = 3;                // Number of pages to allocate (1 page = 4KB)

  // Allocate pages
  Status = gBS->AllocatePages(AllocateAnyPages, EfiBootServicesData, Pages, &Memory);
  if (EFI_ERROR(Status)) {    
    Print(L"TestAllocatePages: AllocatePages failed: %llx\n", Status);
    return Status;
  }

  Print(L"TestAllocatePages: Pages allocated at address: %llx, Pages: %lu\n", Memory, Pages);

  // Use the allocated pages
  SetMem((VOID *)(UINTN)Memory, Pages * 4096, 0);

  // Free the allocated pages
  Status = gBS->FreePages(Memory, Pages);
  if (EFI_ERROR(Status)) {
    Print(L"TestAllocatePages: FreePages failed: %llx\n", Status);
    return Status;
  }

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
UefiMain(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
  TestMemoryAlloc();
  TestAllocateAnyPages();
  TestAllocatePages();
  return EFI_SUCCESS;
}